"""
数据库模型定义
"""
import json
import logging
import mysql.connector
from config import DB_CONFIG
from loguru import logger

class DatabaseConnection:
    """数据库连接管理器"""
    
    @staticmethod
    def connect():
        """建立数据库连接"""
        try:
            connection = mysql.connector.connect(**DB_CONFIG)
            logger.info("数据库连接成功")
            return connection
        except mysql.connector.Error as e:
            logger.error(f"数据库连接失败: {e}")
            raise

class DatabaseSetup:
    """数据库初始化和表创建"""
    
    def __init__(self):
        self.connection = DatabaseConnection.connect()
        self.cursor = self.connection.cursor()
    
    def __del__(self):
        """关闭数据库连接"""
        if hasattr(self, 'cursor') and self.cursor:
            self.cursor.close()
        if hasattr(self, 'connection') and self.connection:
            self.connection.close()
    
    def create_tables(self):
        """创建所需的数据表"""
        try:
            # 创建tasks表
            self.cursor.execute("""
            CREATE TABLE IF NOT EXISTS tasks (
                task_id VARCHAR(36) PRIMARY KEY COMMENT '任务唯一标识',
                entry_url TEXT COMMENT '入口URL地址',
                category VARCHAR(50) COMMENT '任务类别(法规与政策、新闻等)',
                description TEXT COMMENT '任务简述',
                priority VARCHAR(10) COMMENT '信源优先级(高、中、低)',
                source_type VARCHAR(20) COMMENT '信源类型(网页、API等)',
                query_text TEXT COMMENT '用户查询文本',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间'
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='爬虫任务基本信息表';
            """)
            
            # 创建retrieval_conditions表
            self.cursor.execute("""
            CREATE TABLE IF NOT EXISTS retrieval_conditions (
                id INT AUTO_INCREMENT PRIMARY KEY COMMENT '自增ID',
                task_id VARCHAR(36) COMMENT '关联的任务ID',
                keywords JSON COMMENT '关键词列表',
                filters JSON COMMENT '过滤条件',
                date_range JSON COMMENT '日期范围',
                paging_strategy JSON COMMENT '分页策略',
                data_structure_type VARCHAR(50) COMMENT '数据结构类型',
                response_format VARCHAR(20) COMMENT '响应格式',
                data_extraction_path TEXT COMMENT '数据提取路径',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
                FOREIGN KEY (task_id) REFERENCES tasks(task_id) ON DELETE CASCADE
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='检索条件信息表';
            """)
            
            # 创建scraping_strategies表
            self.cursor.execute("""
            CREATE TABLE IF NOT EXISTS scraping_strategies (
                id INT AUTO_INCREMENT PRIMARY KEY COMMENT '自增ID',
                task_id VARCHAR(36) COMMENT '关联的任务ID',
                scraping_frequency VARCHAR(20) COMMENT '爬取频率',
                delay_time INT COMMENT '延迟时间(毫秒)',
                fields_to_extract JSON COMMENT '要提取的字段列表',
                proxy_usage BOOLEAN COMMENT '是否使用代理',
                user_agent_info TEXT COMMENT '用户代理信息',
                anti_crawling_mechanism TEXT COMMENT '反爬机制说明',
                crawling_depth INT COMMENT '爬取深度',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
                FOREIGN KEY (task_id) REFERENCES tasks(task_id) ON DELETE CASCADE
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='爬取策略信息表';
            """)
            
            # 创建authorization_info表
            self.cursor.execute("""
            CREATE TABLE IF NOT EXISTS authorization_info (
                id INT AUTO_INCREMENT PRIMARY KEY COMMENT '自增ID',
                task_id VARCHAR(36) COMMENT '关联的任务ID',
                api_key TEXT COMMENT 'API密钥',
                auth_method VARCHAR(50) COMMENT '授权方式',
                header_info JSON COMMENT '所需头部信息',
                auth_token TEXT COMMENT '认证令牌',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
                FOREIGN KEY (task_id) REFERENCES tasks(task_id) ON DELETE CASCADE
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='授权信息表';
            """)
            
            # 提交事务
            self.connection.commit()
            logger.info("数据表创建成功")
            return True
        except mysql.connector.Error as e:
            self.connection.rollback()
            logger.error(f"创建数据表失败: {e}")
            raise


class DataSourceRepository:
    """数据源仓库，负责数据的存储和检索"""
    
    def __init__(self):
        self.connection = DatabaseConnection.connect()
    
    def __del__(self):
        """关闭数据库连接"""
        if hasattr(self, 'connection') and self.connection:
            self.connection.close()
    
    def save_task(self, task_data):
        """保存任务基本信息"""
        try:
            cursor = self.connection.cursor()
            
            # 检查是否已存在相同query_text的任务
            check_query = "SELECT task_id FROM tasks WHERE query_text = %s"
            cursor.execute(check_query, (task_data.get('query_text'),))
            existing_task = cursor.fetchone()
            
            if existing_task:
                logger.info(f"任务已存在，任务ID: {existing_task[0]}")
                cursor.close()
                return existing_task[0]
            
            # 插入新任务
            insert_query = """
            INSERT INTO tasks (task_id, entry_url, category, description, priority, source_type, query_text)
            VALUES (%s, %s, %s, %s, %s, %s, %s)
            """
            cursor.execute(insert_query, (
                task_data.get('task_id'),
                task_data.get('entry_url'),
                task_data.get('category'),
                task_data.get('description'),
                task_data.get('priority'),
                task_data.get('source_type'),
                task_data.get('query_text')
            ))
            
            self.connection.commit()
            logger.info(f"任务保存成功，任务ID: {task_data.get('task_id')}")
            cursor.close()
            return task_data.get('task_id')
        except mysql.connector.Error as e:
            self.connection.rollback()
            logger.error(f"保存任务失败: {e}")
            raise
    
    def save_retrieval_conditions(self, retrieval_data):
        """保存检索条件信息"""
        try:
            cursor = self.connection.cursor()
            
            # 转换JSON字段
            keywords = json.dumps(retrieval_data.get('keywords', []), ensure_ascii=False)
            filters = json.dumps(retrieval_data.get('filters', {}), ensure_ascii=False)
            date_range = json.dumps(retrieval_data.get('date_range', {}), ensure_ascii=False)
            paging_strategy = json.dumps(retrieval_data.get('paging_strategy', {}), ensure_ascii=False)
            
            insert_query = """
            INSERT INTO retrieval_conditions (
                task_id, keywords, filters, date_range, paging_strategy, 
                data_structure_type, response_format, data_extraction_path
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
            """
            cursor.execute(insert_query, (
                retrieval_data.get('task_id'),
                keywords,
                filters,
                date_range,
                paging_strategy,
                retrieval_data.get('data_structure_type'),
                retrieval_data.get('response_format'),
                retrieval_data.get('data_extraction_path')
            ))
            
            self.connection.commit()
            logger.info(f"检索条件保存成功，任务ID: {retrieval_data.get('task_id')}")
            cursor.close()
            return True
        except mysql.connector.Error as e:
            self.connection.rollback()
            logger.error(f"保存检索条件失败: {e}")
            raise
    
    def save_scraping_strategies(self, strategy_data):
        """保存爬取策略信息"""
        try:
            cursor = self.connection.cursor()
            
            # 转换JSON字段
            fields_to_extract = json.dumps(strategy_data.get('fields_to_extract', []), ensure_ascii=False)
            
            insert_query = """
            INSERT INTO scraping_strategies (
                task_id, scraping_frequency, delay_time, fields_to_extract,
                proxy_usage, user_agent_info, anti_crawling_mechanism, crawling_depth
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
            """
            cursor.execute(insert_query, (
                strategy_data.get('task_id'),
                strategy_data.get('scraping_frequency'),
                strategy_data.get('delay_time'),
                fields_to_extract,
                strategy_data.get('proxy_usage', False),
                strategy_data.get('user_agent_info'),
                strategy_data.get('anti_crawling_mechanism'),
                strategy_data.get('crawling_depth')
            ))
            
            self.connection.commit()
            logger.info(f"爬取策略保存成功，任务ID: {strategy_data.get('task_id')}")
            cursor.close()
            return True
        except mysql.connector.Error as e:
            self.connection.rollback()
            logger.error(f"保存爬取策略失败: {e}")
            raise
    
    def save_authorization_info(self, auth_data):
        """保存授权信息"""
        try:
            cursor = self.connection.cursor()
            
            # 转换JSON字段
            header_info = json.dumps(auth_data.get('header_info', {}), ensure_ascii=False)
            
            insert_query = """
            INSERT INTO authorization_info (
                task_id, api_key, auth_method, header_info, auth_token
            )
            VALUES (%s, %s, %s, %s, %s)
            """
            cursor.execute(insert_query, (
                auth_data.get('task_id'),
                auth_data.get('api_key'),
                auth_data.get('auth_method'),
                header_info,
                auth_data.get('auth_token')
            ))
            
            self.connection.commit()
            logger.info(f"授权信息保存成功，任务ID: {auth_data.get('task_id')}")
            cursor.close()
            return True
        except mysql.connector.Error as e:
            self.connection.rollback()
            logger.error(f"保存授权信息失败: {e}")
            raise
    
    def get_task_by_query(self, query_text):
        """根据查询文本获取任务信息"""
        try:
            cursor = self.connection.cursor(dictionary=True)
            query = "SELECT * FROM tasks WHERE query_text = %s"
            cursor.execute(query, (query_text,))
            result = cursor.fetchone()
            cursor.close()
            return result
        except mysql.connector.Error as e:
            logger.error(f"查询任务失败: {e}")
            raise 